import torch as t
import time
import os


class BasicModule(t.nn.Module):
    """
    封装nn.Module，主要提供save和load两个方法
    """
    def __init__(self):
        super(BasicModule, self).__init__()
        self.module_name = str(type(self))

    def load(self, path, use_gpu=False):
        """
        可加载指定路径的模型
        """
        if not use_gpu:
            self.load_state_dict(t.load(path, map_location=lambda storage, loc: storage))
        else:
            self.load_state_dict(t.load(path))

    def save(self, save_path=None, name=None):
        """
        保存模型，默认使用"模型名字+时间"作为文件名
        """
        prefix = self.module_name + '_'
        if name is None:
            name = time.strftime(prefix + '%m%d_%H-%M-%S.pth')
        if save_path is None:
            save_path = r'./'

        save_path = os.path.join(save_path, prefix + name + time.strftime('%m%d_%H-%M-%S.pth'))
        # 确保保存路径存在
        os.makedirs(os.path.dirname(save_path), exist_ok=True)  # 确保train_log文件夹存在
        t.save(self.state_dict(), save_path)

        return name

    def forward(self, *input):
        pass

